{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "1. 检查是否使用了BF16\n",
    "2. 检查finetune文件中的system与测试函数中的system是否相同"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "591f76d6d8a94fba"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "from transformers import AutoTokenizer\n",
    "from peft import AutoPeftModelForCausalLM\n",
    "import psutil\n",
    "import torch\n",
    "import shutil\n",
    "from modelscope import snapshot_download\n",
    "import subprocess\n",
    "import json\n",
    "import pickle"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "8b8241e99c4a7912"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def gpu_info():\n",
    "    \"\"\"print gpu information\"\"\"\n",
    "\n",
    "    if torch.cuda.is_available():\n",
    "        gpu = torch.cuda.get_device_properties(0)\n",
    "        print(f\"GPU Name: {gpu.name}\")\n",
    "        print(f\"Total GPU RAM: {gpu.total_memory / (1024 ** 3):.2f} GB\")\n",
    "        print(f\"Compute Capability: {gpu.major}.{gpu.minor}\")\n",
    "        print(f\"CUDA Cores: {gpu.multi_processor_count}\")\n",
    "    else:\n",
    "        print(\"No GPU available.\")\n",
    "\n",
    "def system_info(description):\n",
    "    \"\"\"print system information\"\"\"\n",
    "\n",
    "    def system_usage():\n",
    "        # Get system memory information\n",
    "        ram = psutil.virtual_memory()\n",
    "        total_ram = ram.total / (1024 ** 3)\n",
    "        used_ram = ram.used / (1024 ** 3)\n",
    "        ram_percentage = ram.percent\n",
    "\n",
    "        # Get GPU information\n",
    "        if torch.cuda.is_available():\n",
    "            gpu = torch.cuda.get_device_properties(0)\n",
    "            total_gpu_ram = gpu.total_memory / (1024 ** 3)\n",
    "            used_gpu_ram = torch.cuda.memory_allocated(0) / (1024 ** 3)\n",
    "            gpu_percentage = (used_gpu_ram / total_gpu_ram) * 100\n",
    "        else:\n",
    "            total_gpu_ram = 0\n",
    "            used_gpu_ram = 0\n",
    "            gpu_percentage = 0\n",
    "\n",
    "        # Get total disk usage\n",
    "        disk = psutil.disk_usage('/')\n",
    "        total_disk_space = disk.total / (1024 ** 3)\n",
    "        used_disk_space = disk.used / (1024 ** 3)\n",
    "        disk_percentage = disk.percent\n",
    "\n",
    "        return {\n",
    "            \"RAM\": {\n",
    "                \"Total\": total_ram,\n",
    "                \"Used\": used_ram,\n",
    "                \"Percentage\": ram_percentage\n",
    "            },\n",
    "            \"GPU_RAM\": {\n",
    "                \"Total\": total_gpu_ram,\n",
    "                \"Used\": used_gpu_ram,\n",
    "                \"Percentage\": gpu_percentage\n",
    "            },\n",
    "            \"Disk\": {\n",
    "                \"Total\": total_disk_space,\n",
    "                \"Used\": used_disk_space,\n",
    "                \"Percentage\": disk_percentage\n",
    "            }\n",
    "        }\n",
    "\n",
    "    usage_info = system_usage()\n",
    "    print(description)\n",
    "    print(\"System RAM Usage:\")\n",
    "    print(f\"Total RAM: {usage_info['RAM']['Total']} GB\")\n",
    "    print(f\"Used RAM: {usage_info['RAM']['Used']} GB\")\n",
    "    print(f\"RAM Percentage: {usage_info['RAM']['Percentage']} %\")\n",
    "    print()\n",
    "    print(\"GPU RAM Usage:\")\n",
    "    print(f\"Total GPU RAM: {usage_info['GPU_RAM']['Total']} GB\")\n",
    "    print(f\"Used GPU RAM: {usage_info['GPU_RAM']['Used']} GB\")\n",
    "    print(f\"GPU RAM Percentage: {usage_info['GPU_RAM']['Percentage']} %\")\n",
    "    print()\n",
    "    print(\"Disk Usage:\")\n",
    "    print(f\"Total Disk Space: {usage_info['Disk']['Total']} GB\")\n",
    "    print(f\"Used Disk Space: {usage_info['Disk']['Used']} GB\")\n",
    "    print(f\"Disk Percentage: {usage_info['Disk']['Percentage']} %\")"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "33e32158871d40c8"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def get_model(model_id):\n",
    "    \"\"\"从modelscope拉取模型，并将模型移动到model文件夹内\"\"\"\n",
    "\n",
    "    model_dir = snapshot_download(model_id)\n",
    "    shutil.move(model_dir, 'model')\n",
    "    model_path = os.path.abspath(os.path.join('model', os.path.basename(model_dir)))\n",
    "    return model_path"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "31a4c67b151d6d5e"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def get_test(test_data, model, tokenizer):\n",
    "    # 使用测试集数据进行推理\n",
    "    def get_response(row):\n",
    "        response, _ = model.chat(tokenizer, row['question'], history=None, system=\"评估新闻\")\n",
    "        return response\n",
    "    test_data['result_new'] = test_data.apply(get_response, axis=1)\n",
    "    # 重置索引并将字符串的json结果分割成为reason, sentiment, impact三列\n",
    "    test_data['result_new'] = test_data['result_new'].map(lambda x: json.loads(x.replace(\"'\", '\"')))\n",
    "    test_data = test_data.reset_index(drop=True)\n",
    "    test_all = pd.concat([test_data, pd.json_normalize(test_data['result_new'])], axis=1)\n",
    "    # 将新的列后面加_new后缀\n",
    "    rename_dict = {col:col+'_new' for col in ['reason', 'sentiment', 'impact']}\n",
    "    test_all = test_all.rename(columns=rename_dict)\n",
    "    # 进行数据格式转换\n",
    "    test_all[['sentiment_new', 'impact_new']] = test_all[['sentiment_new', 'impact_new']].apply(pd.to_numeric)\n",
    "    # 计算情绪分数的误差\n",
    "    def get_sentiment_error(row):\n",
    "        old = row['sentiment_old']\n",
    "        new = row['sentiment_new']\n",
    "        if old * new >= 0:\n",
    "            error = new - old\n",
    "        else:\n",
    "            error = 1 if abs(new - old) <= 1 else abs(new - old)\n",
    "        return error ** 2\n",
    "    # 计算影响力分数的误差\n",
    "    def get_impact_error(row):\n",
    "        old = row['impact_old']\n",
    "        new = row['impact_new']\n",
    "        return abs(new - old) ** 2\n",
    "    test_all['sentiment_se'] = test_all.apply(get_sentiment_error, axis=1)\n",
    "    test_all['sentiment_se_0.1'] = test_all['sentiment_se'] > 0.1\n",
    "    test_all['impact_se'] = test_all.apply(get_impact_error, axis=1)\n",
    "    test_all['impact_se_0.1'] = test_all['impact_se'] > 0.1\n",
    "    # 打印必要的信息\n",
    "    description = f\"Training epochs: {i+1}\"\n",
    "    system_info(description)\n",
    "    return test_all"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "8e95b65ff0e53c16"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "gpu_info()"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "4ac76fe4d4322611"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "system_info('')"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "aa392ada017588e4"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# 训练文件路径的列表\n",
    "train_path = 'data/train'\n",
    "train_list = [os.path.join(train_path, file) for file in os.listdir(train_path)]\n",
    "# 测试文件的路径\n",
    "test_path = 'data/test'\n",
    "test_list = [pd.read_pickle(os.path.join(test_path, file)) for file in os.listdir(test_path)]"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "9dd6a4c0af4cac19"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "model_path = get_model('TongyiFinance/Tongyi-Finance-14B-Chat')\n",
    "# model_path = ''\n",
    "model_path"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "74a35405d0439fee"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "all_test = []\n",
    "for i, train in enumerate(train_list):\n",
    "    # 设置adapter的文件路径\n",
    "    adapter_path = f\"adapter/{i:02d}\"\n",
    "    # 执行bash文件进行微调\n",
    "    command_bash = f\"bash finetune_lora_single_gpu.sh -m {model_path} -d {train} -o {adapter_path}\"\n",
    "    subprocess.run(command_bash)\n",
    "    # 加载模型\n",
    "    model = AutoPeftModelForCausalLM.from_pretrained(\n",
    "    adapter_path,\n",
    "    device_map=\"auto\",\n",
    "    trust_remote_code=True).eval()\n",
    "    # 加载tokenizer\n",
    "    tokenizer = AutoTokenizer.from_pretrained(adapter_path, trust_remote_code=True)\n",
    "    # 模型效果测试\n",
    "    res_list = [get_test(data, model, tokenizer) for data in test_list]\n",
    "    all_test.append(res_list)\n",
    "    # 合并模型并保存到model文件夹中\n",
    "    ft_model_path = os.path.abspath(f'model/finetune_{i:02d}')\n",
    "    model.merge_and_unload().save_pretrained(ft_model_path, max_shard_size=\"2048MB\", safe_serialization=True)\n",
    "    tokenizer.save_pretrained(ft_model_path)\n",
    "    for file in os.listdir(model_path):\n",
    "        if file.endswith(('.cpp', '.cu')):\n",
    "            source_file_path = os.path.join(model_path, file)\n",
    "            shutil.copy(source_file_path, ft_model_path)\n",
    "    model_path = ft_model_path\n",
    "with open('outcome.pkl', 'wb') as f:\n",
    "    pickle.dump(all_test, f)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "38298803c2ab6437"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
